'''
Description: 情感识别
version: 
Author: 
Date: 2023-12-19 10:30:41
LastEditors: 
LastEditTime: 2023-12-19 10:30:47
'''
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks

from .qwen_api import call

# 'nlp_structbert_sentiment-classification_chinese-base'
#{'scores': [0.9281653165817261, 0.07183466851711273], 'labels': ['负面', '正面']}


# 'nlp_structbert_emotion-classification_chinese-large'
# {'scores': [0.002529235789552331, 0.0031889870297163725, 0.4456832706928253, 0.5213766694068909, 0.0059895701706409454, 0.02061423845589161, 0.0006180171039886773], 'labels': ['惊讶', '高兴', '悲伤', '恐惧', '厌恶', '愤怒', '喜好']}



def _sorted(object,key1='scores',key2='labels'):
  arr1 = object[key1]
  sorted_id = sorted(range(len(arr1)), key=lambda k: arr1[k], reverse=True)
  list1 = []
  list2 = []
  for i in sorted_id:
    list1.append(object[key1][i])
    list2.append(object[key2][i])
  res = {}
  res[key1] = list1
  res[key2] = list2
  return res
  

model_path = ['model/damo/nlp_structbert_emotion-classification_chinese-large','model/damo/nlp_structbert_sentiment-classification_chinese-base']

class SentimentCls:
  #初始化，即默认
  def __init__(self, model_type = 1):
    self.model_type = model_type
    if(model_type == 1 or model_type == 2):
      self.model = pipeline(Tasks.text_classification,model_path[self.model_type-1])
    # if(model_type == 3):
     
  def cls(self,input):
    if(self.model_type == 1 or self.model_type == 2):
      res = _sorted(self.model(input))

      print("\n")
      print(res)
      print("\n")
      ans = res['labels'][0]
      
    if(self.model_type == 3):
       prompt = "你是一个专业的情感分类模型，请根据如下的语句，从[中性、烦恼、乐观、悲伤、激动、紧张、困惑、惊喜、害怕、感激、担心、快乐、放松]中，选择两项类别输出\n" + \
        "样例：\n" + \
        "输入：太好了！\n"+ \
        "输出：快乐、放松\n"+ \
        "输入："+ input +"\n"+ \
        "输出："
       ans = call("qwen-1.8b-chat",prompt)
       res = ans
    
    return res,ans

def main():
  s = SentimentCls(1)
  while(True):
    message=input("输入：")
    res = s.cls(message)
    print("分类：",res)

# if __name__ == '__main__':
#   main()
  




